package org.yanse.utils;

import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

public class TreeListUtil {
    /**
     * @param nodes       要树形化的数据
     * @param getId       获取节点id的方法
     * @param getPid      获取节点父id的方法
     * @param getChildren 获取子集合的方法
     * @param <Node>      节点数据
     * @param <T>         id和pId的数据类型
     * @return 树形化后的数据
     */
    public static <Node, T> List<Node> getTreeList(List<Node> nodes, Function<Node, T> getId, Function<Node, T> getPid, Function<Node, List<Node>> getChildren) {
        Map<T, List<Node>> nodeMap = nodes.stream()
                .filter(node -> getPid.apply(node) != null)
                .collect(Collectors.groupingBy(getPid, Collectors.toList()));

        nodes.forEach(node -> {
            List<Node> nodeList = nodeMap.get(getId.apply(node));
            if (nodeList != null) {
                getChildren.apply(node).addAll(nodeList);
            }
        });
        return nodes;
    }
}
